import tomllib
from pathlib import Path

from pydantic import BaseModel, Field


class DoubaoFlashConfig(BaseModel):
    """DoubaoFlash配置"""

    api_key: str = Field(default="api_key", description="DoubaoFlash API密钥")
    base_url: str = Field(
        default="url",
        description="DoubaoFlash API基础URL",
    )
    model: str = Field(default="model", description="DoubaoFlash模型名称")


class DoubaoProConfig(BaseModel):
    """DoubaoPro配置"""

    api_key: str = Field(default="api_key", description="DoubaoPro API密钥")
    base_url: str = Field(
        default="url",
        description="DoubaoPro API基础URL",
    )
    model: str = Field(default="model", description="DoubaoPro模型名称")


class DoubaoSeedConfig(BaseModel):
    """DoubaoSeed配置"""

    api_key: str = Field(default="api_key", description="DoubaoSeed API密钥")
    base_url: str = Field(
        default="url",
        description="DoubaoSeed API基础URL",
    )
    model: str = Field(default="model", description="DoubaoSeed模型名称")
    enable_thinking: bool = Field(default=True, description="是否启用思考能力")


class EmbeddingConfig(BaseModel):
    """Embedding配置"""

    api_key: str = Field(default="api_key", description="Embedding API密钥")
    base_url: str = Field(
        default="url",
        description="Embedding API基础URL",
    )
    model: str = Field(default="model", description="Embedding模型名称")


def _load_config_from_toml(config_key: str, config_class):
    """通用配置加载函数"""
    config_path = Path(__file__).parent / "settings.toml"

    if config_path.exists():
        try:
            with open(config_path, "rb") as f:
                config_data = tomllib.load(f)
                if config_key in config_data:
                    return config_class(**config_data[config_key])
        except Exception as e:
            print(f"读取配置文件失败: {e}")

    return config_class()


# doubaoFlash配置
def get_doubao_flash_config() -> DoubaoFlashConfig:
    return _load_config_from_toml("doubaoFlash", DoubaoFlashConfig)


# doubaoPro配置
def get_doubao_pro_config() -> DoubaoProConfig:
    return _load_config_from_toml("doubaoPro", DoubaoProConfig)


# doubaoSeed配置
def get_doubao_seed_config() -> DoubaoSeedConfig:
    return _load_config_from_toml("doubaoSeed", DoubaoSeedConfig)


# embedding配置
def get_embedding_config() -> EmbeddingConfig:
    return _load_config_from_toml("embedding", EmbeddingConfig)